# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import gzip
import pickle

from coin.base.datetime_util import to_datetime
from coin.proto import coin_order_enums_pb2
from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.marketsim.core import spec_pb2, marketsim

import experimental.prophet as prophet
from coin.experimental.jaewon.research.adaptive_edge import features, lt_model
from experimental.prophet.ops.linear_model import linear_regression_predict


def _model(symbol_info,
           ref_symbol_nicks,
           model,
           date,
           sampler_offsets=None,
           lt_model_param=None,
           debug=False):
  sampler_offsets = sampler_offsets or {}
  lt_model_param = lt_model_param or {}

  # Timers, sampling points
  ts = prophet.global_timestamp()

  # Symbols
  symbol_nick = symbol_info.nickname
  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs

  ask0_p = book.ask0.price.named_ref('ask0_p')
  bid0_p = book.bid0.price.named_ref('bid0_p')

  # Features
  with prophet.namespace('feature'):
    feature_vars = features.features(ref_symbol_nicks + [symbol_nick], date)

  # Sampler
  with prophet.namespace('sampler'):
    sampler_vars = features.samplers('feature/')
    sampler_vars = {v.name: v for v in sampler_vars}

  # Trigger
  trigger = features.trigger(symbol_nick, 'feature/')

  # Long-term model theo price
  with prophet.namespace('lt_model'):
    theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, date=date, **lt_model_param)

  # Model
  with prophet.namespace('edge_model'):
    x_vars = [v for v in feature_vars if not v.name.endswith('.x')]

    buy_edge_bp = 15.
    sell_edge_bp = 15.

    for sampler_name in model.keys():
      offset = sampler_offsets.get(sampler_name, 3.)

      sampler = sampler_vars[sampler_name]
      buy_model = dict(model[sampler_name]['buy']['model'].params)
      sell_model = dict(model[sampler_name]['sell']['model'].params)
      with prophet.control_if(sampler):
        sampler_buy_edge_bp = linear_regression_predict(x_vars, buy_model) + float(offset)
        sampler_sell_edge_bp = linear_regression_predict(x_vars, sell_model) + float(offset)

      buy_edge_bp = prophet.cond(sampler, sampler_buy_edge_bp, buy_edge_bp)
      sell_edge_bp = prophet.cond(sampler, sampler_sell_edge_bp, sell_edge_bp)

    buy_p = (theo_bid0_p * (1. - buy_edge_bp / 10000.)).named_ref('buy_p')
    sell_p = (theo_ask0_p * (1. + sell_edge_bp / 10000.)).named_ref('sell_p')

  if debug:
    periodic_timer = prophet.timer('1m')
    with prophet.control_if(ts % (60 * 10**9) == 0):
      prophet.printer('%t buy_edge_bp=%s sell_edge_bp=%s', ts, buy_edge_bp, sell_edge_bp)


def _find_order_by_price(orders, price):
  for order in orders:
    if not order.cancel_submitted and order.price == price:
      return order


class Calculator:
  def __init__(self, machine, tick, eval_callback=None, force_passive=False, **kwargs):
    self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick
    self._force_passive = force_passive

    self._graph = prophet.Graph()
    with prophet.as_default_graph(self._graph):
      _model(**kwargs)

    self._ask0_p_var = self._graph.get_named_variable('ask0_p')
    self._bid0_p_var = self._graph.get_named_variable('bid0_p')
    self._buy_p_var = self._graph.get_named_variable('edge_model/buy_p')
    self._sell_p_var = self._graph.get_named_variable('edge_model/sell_p')
    self._timestamp = self._graph.get_named_variable('global_timestamp')
    self._trigger = self._graph.get_named_variable('trigger')

    self._adapter = prophet.FastfeedGraphAdapter(self._graph, machine, eval_callback=eval_callback)

  @property
  def ready(self):
    return not self._buy_p_var.is_null and not self._sell_p_var.is_null

  @property
  def trigger(self):
    return self._trigger.value

  @property
  def timestamp(self):
    return self._timestamp.value

  def calculate_price(self):
    assert self.ready

    if self._force_passive:
      ask0_minus_one = self._tick.get_prev_price(self._ask0_p_var.value)
      bid0_plus_one = self._tick.get_next_price(self._bid0_p_var.value)
      buy_price = min(self._tick.floor(self._buy_p_var.value), ask0_minus_one, bid0_plus_one)
      sell_price = max(self._tick.ceil(self._sell_p_var.value), ask0_minus_one, bid0_plus_one)
    else:
      buy_price = self._tick.floor(self._buy_p_var.value)
      sell_price = self._tick.ceil(self._sell_p_var.value)

    return buy_price, sell_price

  def get_feed_sub_reqs(self, feed_sub_req=None):
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    return self._adapter.get_feed_sub_requests(feed_sub_req)

  @property
  def book_reset_callbacks(self):
    return [self._adapter.on_book_reset]


class PassUnhedgeModelSimStrategy:
  def __init__(self,
               pass_product_sim_spec,
               pass_symbol_info,
               pass_product_tick,
               model_filepath,
               machine,
               date,
               lot_size,
               min_pos,
               max_pos,
               min_fire_period,
               post_only,
               sampler_offsets=None,
               lt_model_param=None,
               sim_spec=None,
               name=None,
               **kwargs):
    self._pass_symbol_info = pass_symbol_info

    self._lot_size = lot_size
    self._min_pos = min_pos
    self._max_pos = max_pos
    self._min_fire_period = min_fire_period
    self._post_only = post_only
    self._name = name

    assert str(pass_symbol_info.product) == pass_product_sim_spec.symbol
    self._msim = marketsim.MarketSim(sim_spec=sim_spec, default_machine=machine)
    self._msim.add_product(pass_symbol_info.product, pass_product_sim_spec)
    self._msim.add_execution_callback(self.on_execution)

    with gzip.open(model_filepath, 'rb') as f:
      model = pickle.load(f)

    self._calc = Calculator(machine=machine,
                            tick=pass_product_tick,
                            eval_callback=self.on_trigger,
                            symbol_info=pass_symbol_info,
                            date=date,
                            model=model,
                            sampler_offsets=sampler_offsets,
                            lt_model_param=lt_model_param,
                            force_passive=post_only,
                            **kwargs)

    self._pos = 0.
    self._last_post_ts = 0
    self._book = None

  @property
  def name(self):
    return self._name

  def manage_buy_order(self, buy_p):
    buy_orders = [
        order for order in self._msim.get_orders()
        if (order.side == coin_order_enums_pb2.BUY_ORDER and order.order_status in (
            spec_pb2.ORDER_STATUS_ON_THE_FLY, spec_pb2.ORDER_STATUS_WORKING))
    ]
    pending_qty = 0.
    for order in buy_orders:
      pending_qty += order.order_qty

    choosen_order = None
    if self._pos + pending_qty < self._max_pos + self._lot_size:
      choosen_order = _find_order_by_price(buy_orders, buy_p)
      if not choosen_order:
        qty = min(self._lot_size, self._max_pos + self._lot_size - self._pos - pending_qty)
        assert qty > 0
        self._msim.submit(product=self._pass_symbol_info.product,
                          side=coin_order_enums_pb2.BUY_ORDER,
                          price=buy_p,
                          qty=qty,
                          post_only=self._post_only)

    for order in buy_orders:
      if (order is not choosen_order and order.order_status != spec_pb2.ORDER_STATUS_DEAD
          and not order.cancel_submitted):
        self._msim.cancel(order_id=order.order_id)

  def manage_sell_order(self, sell_p):
    sell_orders = [
        order for order in self._msim.get_orders()
        if (order.side == coin_order_enums_pb2.SELL_ORDER and order.order_status in (
            spec_pb2.ORDER_STATUS_ON_THE_FLY, spec_pb2.ORDER_STATUS_WORKING))
    ]
    pending_qty = 0.
    for order in sell_orders:
      pending_qty += order.order_qty

    choosen_order = None
    if self._pos + pending_qty < self._max_pos + self._lot_size:
      choosen_order = _find_order_by_price(sell_orders, sell_p)
      if not choosen_order:
        qty = min(self._lot_size, self._max_pos + self._lot_size - self._pos - pending_qty)
        assert qty > 0
        self._msim.submit(product=self._pass_symbol_info.product,
                          side=coin_order_enums_pb2.SELL_ORDER,
                          price=sell_p,
                          qty=qty,
                          post_only=self._post_only)

    for order in sell_orders:
      if (order is not choosen_order and order.order_status != spec_pb2.ORDER_STATUS_DEAD
          and not order.cancel_submitted):
        self._msim.cancel(order_id=order.order_id)

  def on_trigger(self):
    '''
    if not self._calc.ready or not self._calc.trigger or self._book is None:
      return

    ts = self._calc.timestamp
    if ts - self._last_post_ts < self._min_fire_period:
      return

    buy_p, sell_p = self._calc.calculate_price()
    self.manage_buy_order(buy_p)
    self.manage_sell_order(sell_p)
    self._last_post_ts = ts'''
    if not self._calc.ready or not self._calc.trigger or self._book is None:
      return

    ts = self._calc.timestamp
    if ts - self._last_post_ts < self._min_fire_period:
      return

    self._msim.cancel_all()

    buy_p, sell_p = self._calc.calculate_price()
    if self._pos < self._max_pos:
      qty = min(self._max_pos - self._pos, self._lot_size)
      # print('BUY %f %f (%f)' % (buy_p, self._book.bid0().price, (self._book.bid0().price-buy_p)/self._book.bid0().price*10000))
      self._msim.submit(product=self._pass_symbol_info.product,
                        side=coin_order_enums_pb2.BUY_ORDER,
                        price=buy_p,
                        qty=qty,
                        post_only=self._post_only)

    if self._min_pos < self._pos:
      qty = min(self._pos - self._min_pos, self._lot_size)
      # print('SELL %f %f (%f)' % (sell_p, self._book.ask0().price, (sell_p-self._book.ask0().price)/self._book.ask0().price*10000))
      self._msim.submit(product=self._pass_symbol_info.product,
                        side=coin_order_enums_pb2.SELL_ORDER,
                        price=sell_p,
                        qty=qty,
                        post_only=self._post_only)

    self._last_post_ts = ts

  def on_book(self, book):
    self._book = book

  def on_execution(self, execution):
    self._pos = execution.position

  def get_feed_sub_reqs(self, feed_sub_req=None):
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    self._calc.get_feed_sub_reqs(feed_sub_req)
    self._msim.get_feed_sub_reqs(feed_sub_req)
    return feed_sub_req

  def on_book_reset(self, sub_req, book_builder):
    if sub_req.to_str() == self._pass_symbol_info.sub_req:
      book_builder.subscribe_book(self._pass_symbol_info.product, self.on_book)

  @property
  def book_reset_callbacks(self):
    return self._calc.book_reset_callbacks + [self._msim.on_book_reset, self.on_book_reset]

  @property
  def sim(self):
    return self._msim

  def get_sim_stat(self):
    return self._msim.generate_sim_stat()
